import json
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import time

from starlette.responses import StreamingResponse

from LLM_Prompt.prompt_workflow import classify_and_extract
from oracle_api.conn import connect_to_oracle
from boat_nl2sql.utiles import extract_sql_from_markdown
from oracle_api.dispatch_Dataapi import call_api
from ollama_stream.sql_analyse_output_customized import analyze_boat_flow_api_data
from ollama_stream.sql_output import sql_general
from LLM_Prompt.prompt_echarts import visualize_sql_result

app = FastAPI()

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

API_STRATEGIES = [
    "vessel_trajectory", "checkpoint_passage", "violation_detection", "checkpoint_traffic", "basic_information"
]


class QuestionRequest(BaseModel):
    question: str


@app.post("/LLM")
async def query_handler(request: QuestionRequest):
    question = request.question
    if not question:
        raise HTTPException(status_code=400, detail="缺少问题参数")

    # 问题分类+参数提取
    start = time.time()
    classify_prompt = classify_and_extract(question)
    classify_time = time.time() - start
    print(f"问题分类+参数提取花费时间：{classify_time:.4f}秒")
    print(classify_prompt)

    request_dict = json.loads(classify_prompt)
    data = call_api(request_dict)
    action_name, payload = data.get('api', None), data.get('payload', None)

    # 流式响应生成器
    def generate():
        global extracted_sql
        try:
            if action_name in API_STRATEGIES:
                sql_result = data.get('data', None)
                # 构造响应数据
                response_data = {
                    'type': 'result',
                    "result": {
                        "action_name": action_name,
                        "params": payload
                    }
                }
                # 生成SSE格式的响应
                yield f"data: {json.dumps(response_data, ensure_ascii=False)}\n\n"

            else:
                accumulated_content = []
                for content in sql_general(question):
                    print(content, end="", flush=True)  # 外部处理实时输出
                    accumulated_content.append(content)
                    yield f"data: {json.dumps({'type': 'text2sql', 'content': content})}\n\n"

                # 合并所有内容为完整SQL
                full_sql = ''.join(accumulated_content)

                # 提取SQL（假设存在Markdown代码块）
                if "```sql" in full_sql:
                    extracted_sql = extract_sql_from_markdown(full_sql)
                else:
                    extracted_sql = full_sql.strip().strip(';')  # 去除多余分号

                print("sql语句：")
                print(extracted_sql)

                # 执行SQL并获取结果
                try:
                    # 获取数据库连接和游标
                    conn = connect_to_oracle()
                    cursor = conn.cursor()
                    # 执行 SQL 查询，传入参数作为位置绑定
                    cursor.execute(extracted_sql)
                except Exception as e:
                    error_msg = f"SQL执行失败: {str(e)}"
                    yield f"data: {json.dumps({'type': 'error', 'error': error_msg})}\n\n"
                    return  # 终止生成器

                # 获取查询结果的列名
                columns = [col[0] for col in cursor.description]
                sql_result = []
                for row in cursor.fetchall():
                    # 将查询结果转为字典，并转换键值为 Python 原生类型
                    row_dict = dict(zip(columns, row))
                    # 强制将所有键和值转换为 Python 的原生类型（字符串、整数、浮点数等）
                    row_dict = {str(key): (str(value) if value is not None else None) for key, value in
                                row_dict.items()}
                    sql_result.append(row_dict)

                print("查询结果：")
                print(sql_result)

                response_data = {
                    'type': 'result',
                    "result": {
                        "action_name": action_name,
                        "params": payload
                    }
                }

                print("返回result：")
                print(response_data)

                # 生成SSE格式的响应
                yield f"data: {json.dumps(response_data, ensure_ascii=False)}\n\n"

            # sql结果分析
            for content in analyze_boat_flow_api_data(question, sql_result, action_name):
                yield f"data: {json.dumps({'type': 'content', 'content': content})}\n\n"
                print(content, end="", flush=True)  # 外部处理实时输出

            # html图表生成
            if len(sql_result) > 1 and action_name == "not matched":
                try:
                    text2sql_result = []
                    for item in sql_result:
                        name = None
                        count = None
                        for k, v in item.items():
                            if isinstance(v, str):
                                if v.isdigit():
                                    count = int(v)
                                else:
                                    name = v
                        text2sql_result.append((name, count))

                    # 可视化SQL结果
                    output_file = visualize_sql_result(
                        sql_result=text2sql_result,
                        query=extracted_sql,
                        question=question
                    )
                    if output_file:
                        print(f"可视化已完成，请打开以下文件查看: {output_file}")
                        # 读取HTML文件内容
                        try:
                            with open(output_file, 'r', encoding='utf-8') as f:
                                html_content = f.read()
                        except Exception as e:
                            html_content = f"无法读取HTML文件: {str(e)}"

                        yield f"data: {json.dumps({'type': 'html', 'content': html_content})}\n\n"

                except Exception as e:
                    print(f"执行查询或可视化时出错: {e}")

        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'error': str(e)})}\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")


if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=3000)
